#include <lcir/Dag2Dag/D2DDriver.h>
#include <lcir/Dag2Dag/Pattern.h>
#include <lcir/IR/ExprTemps.h>
#include <lcir/IR/Expression.h>
#include <lcir/IR/Op.h>
#include <lcir/Utils/Result.h>

#include <cstddef>
#include <memory>
#include <optional>
#include <vector>

namespace lcir {

static void notifyEraseExpr(Expression *expr, std::vector<Expression *> &exprs) {
  for (auto &ptr : exprs) {
    if (ptr == expr) ptr = nullptr;
  }
}

Pattern *D2DDriver::findPattern(OpInstance *op) {
  auto it = opInstanceToPatterns.find(op);
  if (it == opInstanceToPatterns.end()) return nullptr;
  auto &patterns = it->second;
  int64_t minCostIdx{0};
  size_t idx{0};
  for (auto &p : patterns) {
    if (p->getCost() < patterns[minCostIdx]->getCost()) minCostIdx = idx;
    idx++;
  }
  return patterns[minCostIdx].get();
}

void D2DDriver::addPattern(std::shared_ptr<Pattern> pattern) {
  auto it = opInstanceToPatterns.find(pattern->getRootOp());
  if (it == opInstanceToPatterns.end()) {
    opInstanceToPatterns[pattern->getRootOp()] = std::vector<std::shared_ptr<Pattern>>{};
  }
  opInstanceToPatterns[pattern->getRootOp()].push_back(pattern);

  return;
}

void D2DDriver::run() {
  bool changed = false;
  do {
    changed = false;
    std::vector<Expression *> worklist;
    lambda->walk([&](Expression *expr) { worklist.push_back(expr); });
    while (!worklist.empty()) {
      auto *expr = worklist.back();
      worklist.pop_back();
      if (!expr) continue;
      auto *ret = expr->getResult();
      // 如果发现这玩意使用者为空，直接给他删掉，被特殊标记的不会删(livenode和terminator)
      // TODO: F记得后面把这东西加上！！！！！！！！！！！！！！！
      if (ret && ret->usersEmpty() && !(expr->hasTemp<LiveNode>() || expr->hasTemp<Terminator>())) {
        changed = true;
        notifyEraseExpr(expr, worklist);
        expr->eraseFromLambda();
        continue;
      }
      auto pattern = findPattern(expr->getOpInstance());
      if (!pattern) continue;
      std::vector<Expression *> userExprs;
      for (auto &user : expr->getResult()->getUsers()) {
        userExprs.push_back(user.getExpr());
      }
      auto result = pattern->matchAndRewrite(expr);
      // 匹配成功，更新状态（把他use都丢到worklist里）
      if (succeeded(result)) {
        notifyEraseExpr(expr, worklist);
        // 把userExprs 全部丢到worklist里
        for (auto &user : userExprs) {
          worklist.push_back(user);
        }
        changed = true;
      }
    }
  } while (changed);
}
}  // namespace lcir